#include <immintrin.h>
#include <iostream>
#include <random>
#include "mpi.h"
using namespace std;
#define N 4096

inline bool in(int k, int begin, int end) {
    return k >= begin && k < end;
}

void createbuffer(float* buffer) {
    srand(20000214);
    for (int i = 0; i < N; i++) {
        for (int j = 0; j < N; j++) {
            buffer[i * N + j] = rand() % 100;
            // cout << buffer[i][j] << " ";
        }
        // cout << endl;
    }
    // cout << endl;
}

int main(int argc, char* argv[]) {
    int rank, size;
    double t1, t2;
    float *send_buffer, *recv_buffer;
    float* col = new float[N];
    MPI_Init(&argc, &argv);
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    MPI_Comm_size(MPI_COMM_WORLD, &size);
    int begin, end;
    MPI_Status status;
    begin = rank * (N / size);
    if (rank != size - 1) {
        end = begin + N / size;
    } else {
        end = N;
    }
    if (rank == 0) {
        send_buffer = new float[N * N];
        createbuffer(send_buffer);
        t1 = MPI_Wtime();
    }
    if (rank != size - 1) {
        recv_buffer = new float[N * (N / size)];
    } else {
        recv_buffer = new float[(end - begin) * N];
    }
    for (int i = 0; i < N; i++) {
        MPI_Scatter(&send_buffer[i * N], N / size, MPI_FLOAT,
                    &recv_buffer[i * (N / size)], N / size, MPI_FLOAT, 0,
                    MPI_COMM_WORLD);
        if (N > (N / size) * size) {
            if (rank == 0) {
                MPI_Send(&send_buffer[i * N + (N / size) * size],
                         N - (N / size) * size, MPI_FLOAT, size - 1, size - 1,
                         MPI_COMM_WORLD);
            }
            if (rank == size - 1) {
                MPI_Recv(
                    &recv_buffer[N * (N / size) + i * (end - begin - N / size)],
                    end - begin - N / size, MPI_FLOAT, 0, rank, MPI_COMM_WORLD,
                    &status);
            }
        }
    }

    for (int k = 0; k < N; k++) {
        if (in(k, begin, end)) {
            if (rank == 0) {
                for (int i = k; i < N; i++) {
                    col[i] = send_buffer[i * N + k];
                }
            } else if (rank != size - 1) {
                for (int i = k; i < N; i++) {
                    col[i] = recv_buffer[i * (N / size) + k - begin];
                }
            } else {
                if (k < (N / size) * size) {
                    for (int i = k; i < N; i++) {
                        col[i] = recv_buffer[i * (N / size) + k - begin];
                    }
                } else {
                    for (int i = k; i < N; i++) {
                        col[i] = recv_buffer[N * (N / size) +
                                             i * (N - size * (N / size)) + k -
                                             begin - N / size];
                    }
                }
            }
        }
        int rk = k / (N / size);
        rk = rk >= size ? size - 1 : rk;
        MPI_Bcast(col, N, MPI_FLOAT, rk, MPI_COMM_WORLD);
        for (int i = k; i < N; i++) {
            if (i == k) {
                for (int j = begin; j < end; j++) {
                    if (rank == 0) {
                        send_buffer[i * N + j] /= col[k];
                    } else if (rank != size - 1) {
                        recv_buffer[i * (N / size) + j - begin] /= col[k];
                    } else {
                        if (j < (N / size) * size) {
                            recv_buffer[i * (N / size) + j - begin] /= col[k];
                        } else {
                            cout << j << endl;
                            cout << recv_buffer[N * (N / size) +
                                                (N - size * (N / size)) * i +
                                                j - begin - N / size]
                                 << ' ' << col[k] << endl;
                            recv_buffer[N * (N / size) +
                                        (N - size * (N / size)) * i + j -
                                        begin - N / size] /= col[k];
                        }
                    }
                }

            } else {
                for (int j = begin; j < end; j++) {
                    if (end <= k) {
                        break;
                    }
                    if (j < k) {
                        continue;
                    }
                    if (rank == 0) {
                        if (j != k) {
                            send_buffer[i * N + j] -=
                                send_buffer[i * N + k] * send_buffer[k * N + j];
                        }
                        if (j == end - 1)
                            send_buffer[i * N + k] = 0;
                    } else if (rank != size - 1) {
                        recv_buffer[i * (N / size) + j - begin] -=
                            recv_buffer[k * (N / size) + j - begin] * col[i];
                    } else {
                        if (j < (N / size) * size) {
                            recv_buffer[i * (N / size) + j - begin] -=
                                recv_buffer[k * (N / size) + j - begin] *
                                col[i];
                        } else {
                            recv_buffer[N * (N / size) +
                                        (N - size * (N / size)) * i + j -
                                        begin - N / size] -=
                                recv_buffer[N * (N / size) +
                                            (N - size * (N / size)) * k + j -
                                            begin - N / size] *
                                col[i];
                        }
                    }
                }
            }
        }
        MPI_Barrier(MPI_COMM_WORLD);
    }
    if (rank == 0) {
        for (int i = 0; i < N; i++) {
            for (int j = 0; j < N / size; j++) {
                recv_buffer[i * (N / size) + j] = send_buffer[i * N + j];
            }
        }
    }
    for (int i = 0; i < N; i++) {
        MPI_Gather(&recv_buffer[i * (N / size)], N / size, MPI_FLOAT,
                   &send_buffer[i * N], N / size, MPI_FLOAT, 0, MPI_COMM_WORLD);
        if (N > (N / size) * size) {
            if (rank == 0) {
                MPI_Recv(&send_buffer[i * N + (N / size) * size],
                         N - (N / size) * size, MPI_FLOAT, size - 1, size - 1,
                         MPI_COMM_WORLD, &status);
            }
            if (rank == size - 1) {
                MPI_Send(
                    &recv_buffer[N * (N / size) + i * (end - begin - N / size)],
                    end - begin - N / size, MPI_FLOAT, 0, rank, MPI_COMM_WORLD);
            }
        }
        // for (int j = 0; j < N; j++) {
        //     cout << send_buffer[i * N + j] << ' ';
        // }
        // cout << endl;
    }
    if (rank == 0) {
        t2 = MPI_Wtime();
        cout << "serial: " << (t2 - t1) * 1000 << "ms" << endl;
    }
    MPI_Finalize();
    if (rank == 0)
        delete[] send_buffer;
    delete[] recv_buffer;
    delete[] col;
    return 0;
}